# 要添加一个新单元，输入 '# %%'
# 要添加一个新的标记单元，输入 '# %% [markdown]'
# 我发现写到最后，还是学弟写的那种样子。好麻烦啊，感觉应该一天就能写完，但自己在硬拖。
import copy
import torch
from torch.utils import data
from trainers.fedbase.FedClientBase import FedClientBase

import logging
# logging.basicConfig(level = logging.INFO,format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logging.basicConfig(format ='%(message)s')
logger = logging.getLogger()
logger.setLevel(logging.INFO)

# 训练相关的参数本来应该是服务器发送过来的，包括模型、参数啥的。训练数据应该是本地的。
# 这里为了方便，只模拟联邦学习的效果和逻辑，不模拟联邦学习的真是过程。
# 1. 数据由服务器分成独立同分布，或者非独立同分布。然后分配给各个客户端。而非客户端在本地读取。
# 2. 训练的模型、参数在客户端配置。
# 3. 客户端完成训练过程、服务端完成聚合过程。客户端和服务端只发送和传输模型和优化器的权重参数。完成聚合和分配的通信逻辑。

class FedMetaClient(FedClientBase):
    """训练相关的参数本来应该是服务器发送过来的，包括模型、参数啥的。训练数据应该是本地的。\n
    这里为了方便，只模拟联邦学习的效果和逻辑，不模拟联邦学习的真是过程。
    1. 数据由服务器分成独立同分布，或者非独立同分布。然后分配给各个客户端。而非客户端在本地读取。
    2. 训练的模型、参数在客户端配置。
    3. 客户端完成训练过程、服务端完成聚合过程。客户端和服务端只发送和传输模型和优化器的权重参数。完成聚合和分配的通信逻辑。

    接收参数，完成训练。主要是一次训练的参数。
    每一个算法都应该提供梯度的返回值。即在客户端计算出梯度。其实在服务器计算也无可厚非。

    Args:
        FedClientBase ([type]): [description]

    Returns:
        [type]: [description]
    """    

    def __init__(self,datapair,model):
        super(FedMetaClient, self).__init__(datapair,model)    
        # 服务器传过来的控制参数，和本地生成的控制参数。这个参数直接传聚合就行，不需要求变化。
        self.server_controls =  [torch.zeros_like(param) for param in self.model.parameters()]
        self.client_controls =  [torch.zeros_like(param) for param in self.model.parameters()]
        logger.info("FedMetaClient init----")


    def set_parameters(self,client_params):
        """从服务器加载训练的参数。并根据batch_size等参数划分数据集。
        应该对以下内容进行区分。在联邦个性化层中。包括base_layer和personalization_layer.
        元学习中，并没有将模型进行分层。所以这里不能命名为meta layer

        Args:
            client_params (dict): 客户端进行训练的相关参数
        """     
        # 训练基本参数(包含记载数据)---------------------------------------------------------------------------------------
        super().set_parameters(client_params)   
        # FedAmpClient策略参数-------------------------------------------------------------------------
        # fedper数据集划分参数
        if("support_ratio" in client_params):
            self.support_ratio=client_params["support_ratio"]
        else:
            self.support_ratio = 0.7

        # fedper个性化层的参数
        if("base_layer" in client_params):
            self.base_layer=client_params["base_layer"]
        else:
            self.base_layer = 4

        logger.info("FedMetaClient set_parameters:support_ratio={}\t base_layer={}".format(self.support_ratio,self.base_layer))

    
    def process_data(self):
        return self.process_data_support_and_query()

   
    def process_data_support_and_query(self):
        """考虑到异步梯度下降收敛速度过慢，通信伦次过长。
        对fedmaml的内容进行更改，每一轮进行多次有效梯度下降。

        Returns:
            support_dl,query_dl: 分割好的数据集
        """        
        x,y = self.datapair
        assert len(x) == len(y)
        sq_point  = int(len(y)*self.support_ratio)
        # 将数据转换成tensor,为了加快速度，通过共享内存的方式
        x_support_tensor = torch.FloatTensor(x[:sq_point]).to(self.device)
        y_support_tensor = torch.from_numpy(y[:sq_point]).to(self.device)
        x_query_tensor = torch.FloatTensor(x[sq_point:]).to(self.device)
        y_query_tensor = torch.from_numpy(y[sq_point:]).to(self.device)
        # print(x_train_tensor.dtype,y_train_tensor.dtype)
        support_ds = data.TensorDataset(x_support_tensor,y_support_tensor)
        self.support_dl = data.DataLoader(support_ds,batch_size=self.batch_size,shuffle=True)
        query_ds  = data.TensorDataset(x_query_tensor,y_query_tensor)
        self.query_dl = data.DataLoader(query_ds,batch_size=len(query_ds),shuffle=True)

        self.iteration  = self.epochs*(len(self.support_dl))
        logger.info("FedMetaClient process_data_support_and_query: iteration_support_dl:{}\t".format(self.iteration))
        return self.support_dl,self.query_dl

    # meta load model
    def load_model_state_dict_base_layer(self,model_state_dict):
        # 加载部分参数，还是传递部分参数。传递部分参数过来。加载的时候都是创建一个新的。
        new_par = copy.deepcopy(self.model.state_dict())
        for name in model_state_dict:
            new_par[name]= model_state_dict[name]
        self.model.load_state_dict(new_par)
    

    def get_model_state_dict_base_layer(self):
        state_dict =  self.model.state_dict()
        meta_dict = {key:state_dict[key] for i,key in enumerate(state_dict)  if i < self.base_layer}
        return meta_dict


   
    def train_meta_with_personality_layer(self):
        """当前的训练方式，是基于batch的有效多轮梯度下降。
        也可以实现，基于epoch的有效的多轮梯度下降。

        Returns:
            running_accurcy:训练损失,训练梯度
        """        
        self.model.train()
        loss_sum = 0.
        loss_times = 0.
        correct_sum = 0.
        total_size = 0        
        # 冻结meta层，训练local层。这里可以进行优化，将这样的两个优化器直接放到init中。所有的参数require_grade都是true。但是optimizer只更新部分参数。
        for i,parm in enumerate(self.model.parameters()):
            if(i<self.base_layer):
                parm.requires_grad_(False)
            else:
                parm.requires_grad=True
        self.optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, self.model.parameters()),lr=0.01)

        for epoch in range(self.epochs):
            # 使用support_set进行梯度下降
            for batch_x,batch_y in self.support_dl:
                # 正向传播
                outputs = self.model(batch_x)

                # 计算loss
                loss = self.loss_func(outputs,batch_y.long())
                loss_sum +=loss.item()
                loss_times+=1
                # 反向传播
                self.optimizer.zero_grad()
                loss.backward()

                # 梯度下降
                self.optimizer.step()
            
                # 开启meta层，冻结local层。使用query_set进行一次梯度下降
                for i,parm in enumerate(self.model.parameters()):
                    if(i<self.base_layer):
                        parm.requires_grad=True
                    else:
                        parm.requires_grad=False
                self.optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, self.model.parameters()),lr=0.01)

                for batch_x,batch_y in self.query_dl:
                    # 正向传播
                    outputs = self.model(batch_x)
                    _,predicted = torch.max(outputs,1)
                    correct_sum += (predicted == batch_y).sum().item()
                    total_size += len(batch_y)
                    # 计算loss
                    loss = self.loss_func(outputs,batch_y.long())
                    loss_sum +=loss.item()
                    loss_times+=1
                    # 反向传播
                    self.optimizer.zero_grad()
                    loss.backward()

                    # 梯度下降
                    self.optimizer.step()
                # print(self.optimizer.state_dict())
                # 解冻所有的层
                for parm in self.model.parameters():
                    parm.requires_grad=True
                self.optimizer = torch.optim.SGD(self.model.parameters(),lr=0.01)

        return loss_sum/loss_times,100*correct_sum/total_size
    
    # 元数据，首先进行本地的fine_tuning
    def test_meta_with_personality_layer(self):
        test_model = copy.deepcopy(self.model)
        test_model.train()
        # 冻结meta层，训练local层。这里可以进行优化，将这样的两个优化器直接放到init中。所有的参数require_grade都是true。但是optimizer只更新部分参数。
        for i,parm in enumerate(test_model.parameters()):
            if(i<self.base_layer):
                parm.requires_grad_(False)
            else:
                parm.requires_grad=True
        test_optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, test_model.parameters()),lr=self.inner_lr )

        for epoch in range(self.epochs):
            # 使用support_set进行梯度下降
            for batch_x,batch_y in self.support_dl:
                # 正向传播
                outputs = test_model(batch_x)

                # 计算loss
                loss = self.loss_func(outputs,batch_y.long())
                # 反向传播
                test_optimizer.zero_grad()
                loss.backward()

                # 梯度下降
                test_optimizer.step()
        # 解冻所有的层
        for parm in test_model.parameters():
            parm.requires_grad=True
        test_optimizer = torch.optim.SGD(test_model.parameters(),lr=self.inner_lr)

        test_model.eval()
        correct = 0
        total =0
        with torch.no_grad():
            for features,labels in self.query_dl:
                outputs = test_model(features)
                # print(outputs.shape)
                _,predicted = torch.max(outputs,1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        accuracy = 100*correct/total
        return accuracy




    # 用来训练当前的本地模型。每次训练返回损失值running_loss和running_accuracy
    def train_meta_with_multi_maml(self):
        self.model.train()
        loss_sum = 0.
        loss_times = 0.
        correct_sum = 0.
        total_size = 0        
        # 保留原始梯度的一个副本
        old_model_state_dict = copy.deepcopy(self.model.state_dict())
        for epoch in range(self.epochs):
            # 使用support_set进行梯度下降
            for batch_x,batch_y in self.support_dl:
                # 正向传播
                outputs = self.model(batch_x)

                # 计算loss
                loss = self.loss_func(outputs,batch_y.long())
                loss_sum +=loss.item()
                loss_times+=1
                # 反向传播
                self.optimizer.zero_grad()
                loss.backward()

                # 梯度下降
                self.optimizer.step()
                    
                # 使用query_set进行一次梯度下降
                for batch_x,batch_y in self.query_dl:
                    # 正向传播
                    outputs = self.model(batch_x)
                    _,predicted = torch.max(outputs,1)
                    correct_sum += (predicted == batch_y).sum().item()
                    total_size += len(batch_y)
                    # 计算loss
                    loss = self.loss_func(outputs,batch_y.long())
                    loss_sum +=loss.item()
                    loss_times+=1
                    # 反向传播
                    self.optimizer.zero_grad()
                    loss.backward()

                    # 梯度下降，在下降前加载初始的参数状态
                    self.model.load_state_dict(old_model_state_dict)
                    self.optimizer.step()
                    # 需要将新的模型参数保存，作为下一次queryset梯度下降的起始点
                    old_model_state_dict = copy.deepcopy(self.model.state_dict())
        # print(self.optimizer.state_dict())
        return loss_sum/loss_times,100*correct_sum/total_size


    # 进行finetuning之后测试
    def test_meta_with_multi_maml(self):
        test_model = copy.deepcopy(self.model)
        test_optimizer = torch.optim.SGD(test_model.parameters(),lr=self.inner_lr)
        test_model.train()
        # 冻结meta层，训练local层。这里可以进行优化，将这样的两个优化器直接放到init中。所有的参数require_grade都是true。但是optimizer只更新部分参数。
        for epoch in range(self.epochs):
            # 使用support_set进行梯度下降
            for batch_x,batch_y in self.support_dl:
                # 正向传播
                outputs = test_model(batch_x)

                # 计算loss
                loss = self.loss_func(outputs,batch_y.long())
                # 反向传播
                test_optimizer.zero_grad()
                loss.backward()

                # 梯度下降
                test_optimizer.step()
        # 进入测试
        test_model.eval()
        correct = 0
        total =0
        with torch.no_grad():
            for features,labels in self.query_dl:
                outputs = test_model(features)
                # print(outputs.shape)
                _,predicted = torch.max(outputs,1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        accuracy = 100*correct/total
        return accuracy

 
    # 用来训练当前的本地模型。每次训练返回损失值running_loss和running_accuracy
    def train_meta_with_multi_reptile(self):
        self.model.train()
        loss_sum = 0.
        loss_times = 0.
        correct_sum = 0.
        total_size = 0        
        # 保留原始梯度的一个副本
        model_state_dict = copy.deepcopy(self.model.state_dict())
        for epoch in range(self.epochs):
            # 使用support_set和queryset进行梯度下降
            for dl in (self.support_dl,self.query_dl):
                for batch_x,batch_y in dl:
                    # 正向传播
                    outputs = self.model(batch_x)
                    _,predicted = torch.max(outputs,1)
                    correct_sum += (predicted == batch_y).sum().item()
                    total_size += len(batch_y)
                    
                    # 计算loss
                    loss = self.loss_func(outputs,batch_y.long())
                    loss_sum +=loss.item()
                    loss_times+=1
                    # 反向传播
                    self.optimizer.zero_grad()
                    loss.backward()

                    # 梯度下降
                    self.optimizer.step()
            
        # reptile不需要进行queryset的梯度下降。而是使用outer_lr进行放缩。这个放缩与maml一样，放到服务器上。

        # print(self.optimizer.state_dict())
        return loss_sum/loss_times,100*correct_sum/total_size


    # 进行finetuning之后测试
    def test_meta_with_multi_reptile(self):
        test_model = copy.deepcopy(self.model)
        test_optimizer = torch.optim.SGD(test_model.parameters(),lr=self.inner_lr)
        test_model.train()
        # 冻结meta层，训练local层。这里可以进行优化，将这样的两个优化器直接放到init中。所有的参数require_grade都是true。但是optimizer只更新部分参数。
        for epoch in range(self.epochs):
            # 使用support_set进行梯度下降
            for batch_x,batch_y in self.support_dl:
                # 正向传播
                outputs = test_model(batch_x)

                # 计算loss
                loss = self.loss_func(outputs,batch_y.long())
                # 反向传播
                test_optimizer.zero_grad()
                loss.backward()

                # 梯度下降
                test_optimizer.step()
        # 进入测试
        test_model.eval()
        correct = 0
        total =0
        with torch.no_grad():
            for features,labels in self.query_dl:
                outputs = test_model(features)
                # print(outputs.shape)
                _,predicted = torch.max(outputs,1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        accuracy = 100*correct/total
        return accuracy




    # 加载控制参数
    def load_controls(self,controls_list):
        self.server_controls = copy.deepcopy(controls_list)

    # 返回控制参数
    def get_controls(self):
        return copy.deepcopy(self.client_controls)
    
    # 用来训练当前的本地模型。每次训练返回损失值running_loss和running_accuracy
    def train_meta_with_reptile_adjust(self):
        self.model.train()
        loss_sum = 0.
        loss_times = 0.
        correct_sum = 0.
        total_size = 0        
        # 保留一份原始参数
        old_param_list = [torch.clone(parm).detach() for parm in self.model.parameters()]
        # 计算两个不同的差异。用来衡量两个分布的差异程度。
        # 分布的差异越大，这个修正的服务应该越小，防止出现，过大修正，导致最后的发现。
        # 当分布的差异接近与1 的时候，进行全额的修正。因为太大的修正会导致最后本地模型中。该修正锁占的影响过大，最后出现发散。
        # 排除因为调整因子，起到过大的作用，导致每次更新波动太大，没办法收敛。
        server_client_controls_variance = self.get_params_norm(self.server_controls,self.client_controls)
        adaption = 1
        if server_client_controls_variance < 2 :
            adaption=1
        else:
            adaption = server_client_controls_variance
        
        for epoch in range(self.epochs):
            # 使用support_set和queryset进行梯度下降
            for dl in (self.support_dl,self.query_dl):
                for batch_x,batch_y in dl:
                    # 正向传播，计算准确度
                    outputs = self.model(batch_x)
                    # 计算loss
                    loss = self.loss_func(outputs,batch_y.long())
                    # 反向传播
                    self.optimizer.zero_grad()
                    loss.backward()
                    # 梯度下降，修改模型梯度。将当前的控制参数放到梯度当中之后在修改。防止破坏现有的优化器，继续使用原来的方法进行梯度下降。
                    for param, c, ci in zip(self.model.parameters(), self.server_controls, self.client_controls):
                        if param is None:
                            continue
                        # 修改现有的参数的grad。这相当与添加一个更新防线，与服务器上的更新方向一致。
                        # 我终于知道weight_control的意思了。他妈的。应该是用来平衡这里的更新次数和学习率。
                        # 使得经过K次更新后，正好累加了一次服务器上参数的更新方向。
                        param.grad.add_((c - ci)/adaption)
                    self.optimizer.step()

                    _,predicted = torch.max(outputs,1)
                    correct_sum += (predicted == batch_y).sum().item()
                    total_size += len(batch_y)
                    loss_sum +=loss.item()
                    loss_times+=1
            
        # reptile不需要进行queryset的梯度下降。而是使用outer_lr进行放缩。这个放缩与maml一样，放到服务器上。
        
        # 计算self.delta_model模型变化
        new_param_list = [torch.clone(parm).detach() for parm in self.model.parameters()]
        delta_model = []
        for new_param, old_param in zip(new_param_list, old_param_list):
            delta_model.append(new_param - old_param)
        
        old_controls = copy.deepcopy(self.client_controls)
        
        # 得到本地的控制参数。这个控制参数，表示，与服务器的控制参数的差，减去，变化的差。表示变化的剧烈程度
        # 相当于对本地的变化起到一定的抑制作用。待会执行以下，看看这个东西能不能，真的进行赋值。
        # 计算new_controls控制参数。我觉得这个有点多余。并不多余，非常有用
        weight_control = 1 /(self.inner_lr* loss_times)
        # 更新delta_controls和本地client_controls
        strategy = 1
        # 每一次本地更新的方向。如果将一下指放到grad中，相当于每次添加一个向服务器的方向更新的方向。
        if strategy==1:
            for i in range(len(self.client_controls)):
                self.client_controls[i] = -delta_model[i] * weight_control
        else:
            # 这玩意表示模型负变化的方向。也就是本地梯度的方向。但又考虑到了历史的梯度。
            # 在梯度下降的时候add_(-ci)表示减去上次的梯度值，抵消掉上次的梯度变化。这，好像又缺乏意义。
            # 是不是应该给个参数，减小一下影响更好。加一个参数0.2，用来微调 ，而不是真实影响。
            # 让历史的影响降低一些。之前的方向保留0.2.最新的方向保留0.8。
            # 在范数保持一致。事实上，原来的范数就是一致的。
            # 另外考虑到loss溢出的问题，估计多半是因为本地更新方向和服务器更新方向相差太多导致的。
            # 可以通过计算两个control的方法，显示这种差异。
            # 猜测出现这种情况的原因是，因为两个方向完全不一致，并且由于batch过多导致，更新次数太多，加大了这种差异。
            # 最终根本没办法收敛。所以，大胆猜测，可以减小本地更新次数。
            for i in range(len(self.client_controls)):
                self.client_controls[i] =-(self.server_controls[i]-self.client_controls[i] + delta_model[i] * weight_control)
                
        
        # 计算模型的共享两
        model_variance = self.get_params_norm(old_param_list,new_param_list)
        control_variance = self.get_params_norm(old_controls,self.client_controls)


        # print(self.optimizer.state_dict())
        return loss_sum/loss_times,100*correct_sum/total_size,server_client_controls_variance,model_variance,control_variance



    # 进行finetuning之后测试
    def test_meta_with_reptile_adjust(self):
        test_model = copy.deepcopy(self.model)
        test_optimizer = torch.optim.SGD(test_model.parameters(),lr=self.inner_lr)
        test_model.train()
        # 冻结meta层，训练local层。这里可以进行优化，将这样的两个优化器直接放到init中。所有的参数require_grade都是true。但是optimizer只更新部分参数。
        for epoch in range(self.epochs):
            # 使用support_set进行梯度下降
            for batch_x,batch_y in self.support_dl:
                # 正向传播
                outputs = test_model(batch_x)

                # 计算loss
                loss = self.loss_func(outputs,batch_y.long())
                # 反向传播
                test_optimizer.zero_grad()
                loss.backward()

                # 梯度下降
                test_optimizer.step()
        # 进入测试
        test_model.eval()
        correct = 0
        total =0
        with torch.no_grad():
            for features,labels in self.query_dl:
                outputs = test_model(features)
                # print(outputs.shape)
                _,predicted = torch.max(outputs,1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        accuracy = 100*correct/total
        return accuracy

